#include "market_persistor.h"

#include <iostream>
#include <chrono>
#include <filesystem>
#include <iomanip>
#include <sstream>
#include <arrow/api.h>
#include <arrow/io/api.h>
#include <parquet/arrow/writer.h>

MarketPersistor::MarketPersistor(MARKETStruct* market_ptr)
    : market_ptr_(market_ptr) {}

MarketPersistor::~MarketPersistor() {
    stop();
}

void MarketPersistor::start() {
    if (running_) {
        std::cerr << "⚠️ MarketPersistor 已经在运行\n";
        return;
    }

    running_ = true;
    worker_ = std::thread([this]() { this->persist_loop(); });
    std::cout << "✅ MarketPersistor 启动成功\n";
}

void MarketPersistor::stop() {
    if (!running_) return;
    running_ = false;
    if (worker_.joinable()) worker_.join();
    std::cout << "🛑 MarketPersistor 已停止\n";
}

void MarketPersistor::persist_loop() {
    while (running_) {
        try {
            write_price_level();
            write_order();
            write_trade();
            write_order_rsp();
        } catch (const std::exception& e) {
            std::cerr << "❌ 写入 parquet 失败: " << e.what() << std::endl;
        }

        std::this_thread::sleep_for(std::chrono::seconds(1));
    }
}

// ✅ 自动生成 "../data/yyyyMMdd/<name>.parquet"
std::string MarketPersistor::get_output_path(const std::string& name) const {
    auto now = std::chrono::system_clock::now();
    std::time_t t = std::chrono::system_clock::to_time_t(now);
    std::tm tm{};
    localtime_r(&t, &tm);

    std::ostringstream oss;
    oss << "../data/"
        << std::put_time(&tm, "%Y%m%d");

    std::filesystem::create_directories(oss.str());  // 确保目录存在
    oss << "/" << name << ".parquet";
    return oss.str();
}

// =============================
// 写 price_level.parquet（增量）
// =============================
void MarketPersistor::write_price_level() {
    auto& shm = market_ptr_->priceLevelShm;
    size_t write_index = shm.next_write.load(std::memory_order_relaxed);

    if (write_index == last_price_level_index_) return;  // 没新数据

    arrow::MemoryPool* pool = arrow::default_memory_pool();
    arrow::UInt8Builder exchange(pool);
    arrow::UInt32Builder market(pool);
    arrow::Int64Builder price(pool);
    arrow::DoubleBuilder size(pool);
    arrow::UInt8Builder side(pool);

    while (last_price_level_index_ < write_index) {
        const auto& rec = shm[last_price_level_index_ % shm._size];
        exchange.Append(rec.exchange);
        market.Append(rec.market_id);
        price.Append(rec.pl.price);
        size.Append(rec.pl.size);
        side.Append(static_cast<uint8_t>(rec.order_side));
        ++last_price_level_index_;
    }

    std::shared_ptr<arrow::Array> a,b,c,d,e;
    exchange.Finish(&a);
    market.Finish(&b);
    price.Finish(&c);
    size.Finish(&d);
    side.Finish(&e);

    auto schema = arrow::schema({
        arrow::field("exchange", arrow::uint8()),
        arrow::field("market_id", arrow::uint32()),
        arrow::field("price", arrow::int64()),
        arrow::field("size", arrow::float64()),
        arrow::field("side", arrow::uint8())
    });

    auto table = arrow::Table::Make(schema, {a,b,c,d,e});
    auto path = get_output_path("price_level");

    auto fout = arrow::io::FileOutputStream::Open(path, true);
    if (!fout.ok()) throw std::runtime_error(fout.status().ToString());

    parquet::WriterProperties::Builder builder;
    builder.compression(parquet::Compression::SNAPPY);
    auto st = parquet::arrow::WriteTable(*table, pool, *fout, 1024, builder.build());
    if (!st.ok()) throw std::runtime_error(st.ToString());

    std::cout << "📈 写入 price_level " << (write_index - last_price_level_index_) << " 条\n";
}

// =============================
// 写 order.parquet（增量）
// =============================
void MarketPersistor::write_order() {
    auto& shm = market_ptr_->orderShm;
    size_t write_index = shm.next_write.load(std::memory_order_relaxed);

    if (write_index == last_order_index_) return;

    arrow::MemoryPool* pool = arrow::default_memory_pool();
    arrow::UInt8Builder exchange(pool);
    arrow::UInt32Builder market(pool);
    arrow::UInt8Builder action(pool);
    arrow::UInt64Builder order_id(pool);
    arrow::UInt8Builder side(pool);
    arrow::DoubleBuilder size(pool);
    arrow::Int64Builder price(pool);

    while (last_order_index_ < write_index) {
        const auto& rec = shm[last_order_index_ % shm._size];
        exchange.Append(rec.exchange);
        market.Append(rec.market_id);
        action.Append(static_cast<uint8_t>(rec.order_type));
        order_id.Append(rec.order_id);
        side.Append(static_cast<uint8_t>(rec.order_side));
        size.Append(rec.size);
        price.Append(rec.price);
        ++last_order_index_;
    }

    std::shared_ptr<arrow::Array> a,b,c,d,e,f,g;
    exchange.Finish(&a);
    market.Finish(&b);
    action.Finish(&c);
    order_id.Finish(&d);
    side.Finish(&e);
    size.Finish(&f);
    price.Finish(&g);

    auto schema = arrow::schema({
        arrow::field("exchange", arrow::uint8()),
        arrow::field("market_id", arrow::uint32()),
        arrow::field("order_type", arrow::uint8()),
        arrow::field("order_id", arrow::uint64()),
        arrow::field("order_side", arrow::uint8()),
        arrow::field("size", arrow::float64()),
        arrow::field("price", arrow::int64())
    });

    auto table = arrow::Table::Make(schema, {a,b,c,d,e,f,g});
    auto path = get_output_path("order");

    auto fout = arrow::io::FileOutputStream::Open(path, true);
    if (!fout.ok()) throw std::runtime_error(fout.status().ToString());

    parquet::WriterProperties::Builder builder;
    builder.compression(parquet::Compression::SNAPPY);
    auto st = parquet::arrow::WriteTable(*table, pool, *fout, 1024, builder.build());
    if (!st.ok()) throw std::runtime_error(st.ToString());

    std::cout << "📦 写入 order " << (write_index - last_order_index_) << " 条\n";
}

// =============================
// 写 trade.parquet（增量）
// =============================
void MarketPersistor::write_trade() {
    auto& shm = market_ptr_->tradeShm;
    size_t write_index = shm.next_write.load(std::memory_order_relaxed);

    if (write_index == last_trade_index_) return;

    arrow::MemoryPool* pool = arrow::default_memory_pool();
    arrow::UInt8Builder exchange(pool);
    arrow::UInt32Builder market(pool);
    arrow::UInt64Builder trade_id(pool);
    arrow::DoubleBuilder size(pool);
    arrow::Int64Builder price(pool);
    arrow::UInt64Builder ask_id(pool);
    arrow::UInt64Builder bid_id(pool);
    arrow::UInt64Builder ts(pool);

    while (last_trade_index_ < write_index) {
        const auto& rec = shm[last_trade_index_ % shm._size];
        exchange.Append(rec.exchange);
        market.Append(rec.market_id);
        trade_id.Append(rec.trade_id);
        size.Append(rec.size);
        price.Append(rec.price);
        ask_id.Append(rec.ask_id);
        bid_id.Append(rec.bid_id);
        ts.Append(rec.timestamp);
        ++last_trade_index_;
    }

    std::shared_ptr<arrow::Array> a,b,c,d,e,f,g,h;
    exchange.Finish(&a);
    market.Finish(&b);
    trade_id.Finish(&c);
    size.Finish(&d);
    price.Finish(&e);
    ask_id.Finish(&f);
    bid_id.Finish(&g);
    ts.Finish(&h);

    auto schema = arrow::schema({
        arrow::field("exchange", arrow::uint8()),
        arrow::field("market_id", arrow::uint32()),
        arrow::field("trade_id", arrow::uint64()),
        arrow::field("size", arrow::float64()),
        arrow::field("price", arrow::int64()),
        arrow::field("ask_id", arrow::uint64()),
        arrow::field("bid_id", arrow::uint64()),
        arrow::field("timestamp", arrow::uint64())
    });

    auto table = arrow::Table::Make(schema, {a,b,c,d,e,f,g,h});
    auto path = get_output_path("trade");

    auto fout = arrow::io::FileOutputStream::Open(path, true);
    if (!fout.ok()) throw std::runtime_error(fout.status().ToString());

    parquet::WriterProperties::Builder builder;
    builder.compression(parquet::Compression::SNAPPY);
    auto st = parquet::arrow::WriteTable(*table, pool, *fout, 1024, builder.build());
    if (!st.ok()) throw std::runtime_error(st.ToString());

    std::cout << "💹 写入 trade " << (write_index - last_trade_index_) << " 条\n";
}

// =============================
// 写 order_rsp.parquet（增量）
// =============================
void MarketPersistor::write_order_rsp() {
    auto& shm = market_ptr_->orderRspShm;
    size_t write_index = shm.next_write.load(std::memory_order_relaxed);

    if (write_index == last_order_rsp_index_) return;

    arrow::MemoryPool* pool = arrow::default_memory_pool();
    arrow::UInt8Builder exchange(pool);
    arrow::UInt32Builder market(pool);
    arrow::UInt64Builder client_order_id(pool);
    arrow::UInt64Builder order_id(pool);
    arrow::DoubleBuilder size(pool);
    arrow::Int64Builder price(pool);
    arrow::UInt8Builder status(pool);
    arrow::UInt64Builder send_ts(pool);
    arrow::UInt64Builder recv_ts(pool);

    while (last_order_rsp_index_ < write_index) {
        const auto& rec = shm[last_order_rsp_index_ % shm._size];
        exchange.Append(rec.exchange);
        market.Append(rec.market_id);
        client_order_id.Append(rec.client_order_id);
        order_id.Append(rec.order_id);
        size.Append(rec.size);
        price.Append(rec.price);
        status.Append(static_cast<uint8_t>(rec.status));
        send_ts.Append(rec.send_timestamp);
        recv_ts.Append(rec.receive_timestamp);
        ++last_order_rsp_index_;
    }

    std::shared_ptr<arrow::Array> a,b,c,d,e,f,g,h,i;
    exchange.Finish(&a);
    market.Finish(&b);
    client_order_id.Finish(&c);
    order_id.Finish(&d);
    size.Finish(&e);
    price.Finish(&f);
    status.Finish(&g);
    send_ts.Finish(&h);
    recv_ts.Finish(&i);

    auto schema = arrow::schema({
        arrow::field("exchange", arrow::uint8()),
        arrow::field("market_id", arrow::uint32()),
        arrow::field("client_order_id", arrow::uint64()),
        arrow::field("order_id", arrow::uint64()),
        arrow::field("size", arrow::float64()),
        arrow::field("price", arrow::int64()),
        arrow::field("status", arrow::uint8()),
        arrow::field("send_timestamp", arrow::uint64()),
        arrow::field("receive_timestamp", arrow::uint64())
    });

    auto table = arrow::Table::Make(schema, {a,b,c,d,e,f,g,h,i});
    auto path = get_output_path("order_rsp");

    auto fout = arrow::io::FileOutputStream::Open(path, true);
    if (!fout.ok()) throw std::runtime_error(fout.status().ToString());

    parquet::WriterProperties::Builder builder;
    builder.compression(parquet::Compression::SNAPPY);
    auto st = parquet::arrow::WriteTable(*table, pool, *fout, 1024, builder.build());
    if (!st.ok()) throw std::runtime_error(st.ToString());

    std::cout << "📬 写入 order_rsp " << (write_index - last_order_rsp_index_) << " 条\n";
}
